#include <iostream>
#include <fcntl.h>
#include <unistd.h>
#include <termios.h>
#include <cstring>
#include <vector>
#include <cstdio>
#include <stdexcept>
#include <memory>
#include <chrono>
#include <thread>

#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <sensor_msgs/msg/magnetic_field.hpp>
#include <geometry_msgs/msg/vector3.hpp>
#include <std_msgs/msg/header.hpp>

// 极简串口封装（非阻塞读），C++11
class ImuSensor {
public:
    ImuSensor(const std::string& port, int baud) : uart_port(port), baudrate(baud), uart_fd(-1) {}

    bool init() {
        uart_fd = open(uart_port.c_str(), O_RDWR | O_NOCTTY | O_NONBLOCK);
        if (uart_fd == -1) {
            perror("打开串口失败");
            return false;
        }

        struct termios options;
        if (tcgetattr(uart_fd, &options) != 0) {
            perror("获取串口属性失败");
            close(uart_fd);
            uart_fd = -1;
            return false;
        }

        speed_t sp = B115200;
        if (baudrate == 9600) sp = B9600;
        else if (baudrate == 19200) sp = B19200;
        else if (baudrate == 38400) sp = B38400;
        else if (baudrate == 57600) sp = B57600;
        else if (baudrate == 115200) sp = B115200;
        else if (baudrate == 230400) sp = B230400;
#ifdef B460800
        else if (baudrate == 460800) sp = B460800;
#endif
#ifdef B921600
        else if (baudrate == 921600) sp = B921600;
#endif

        cfsetispeed(&options, sp);
        cfsetospeed(&options, sp);

        options.c_cflag &= ~CSIZE;
        options.c_cflag |= CS8;       // 8N1
        options.c_cflag &= ~PARENB;
        options.c_cflag &= ~CSTOPB;
        options.c_cflag |= (CLOCAL | CREAD);

        options.c_iflag = 0;
        options.c_oflag = 0;
        options.c_lflag = 0;          // 原始模式
        options.c_cc[VMIN]  = 0;      // 非阻塞模式
        options.c_cc[VTIME] = 0;      // 无超时

        if (tcsetattr(uart_fd, TCSANOW, &options) != 0) {
            perror("设置串口属性失败");
            close(uart_fd);
            uart_fd = -1;
            return false;
        }

        tcflush(uart_fd, TCIOFLUSH);
        return true;
    }

    ssize_t readSome(uint8_t* buf, size_t cap) {
        if (uart_fd == -1) return -1;
        ssize_t result = read(uart_fd, buf, cap);
        if (result == -1 && errno == EAGAIN) {
            return 0; // 非阻塞模式下没有数据可读
        }
        return result;
    }

    ~ImuSensor() {
        if (uart_fd != -1) {
            close(uart_fd);
            uart_fd = -1;
        }
    }

private:
    std::string uart_port;
    int baudrate;
    int uart_fd;
};

// XBUS 工具函数
static inline uint16_t be16(const uint8_t* p){ return (uint16_t(p[0])<<8)|p[1]; }
static inline uint32_t be32(const uint8_t* p){ return (uint32_t(p[0])<<24)|(uint32_t(p[1])<<16)|(uint32_t(p[2])<<8)|p[3]; }
static inline float befloat32(const uint8_t* p){ uint32_t u = be32(p); float v; std::memcpy(&v, &u, 4); return v; }

static inline uint8_t checksum_basic(uint8_t busId, uint8_t mid, const uint8_t* payload, size_t n, int extLenBytes){
    // 校验和包含 BusID + MID + LEN(或0xFF+LENEX_H+LENEX_L) + PAYLOAD
    uint32_t s = 0;
    s += busId;
    s += mid;
    if (extLenBytes == 0) {
        s += (uint8_t)n;
    } else {
        s += 0xFF;
        s += ((n >> 8) & 0xFF);
        s += (n & 0xFF);
    }
    for (size_t i=0;i<n;++i) s += payload[i];
    return (uint8_t)((0x100 - (s & 0xFF)) & 0xFF);
}

// 遍历 MTData2 的 Packet
template<typename F>
static bool foreachPacket(const std::vector<uint8_t>& pl, F&& cb){
    size_t i = 0, n = pl.size();
    while (i + 3 <= n) {
        uint16_t did = be16(&pl[i]);
        uint8_t sz = pl[i+2];
        i += 3;
        if (i + sz > n) return false; // 截断
        cb(did, &pl[i], sz);
        i += sz;
    }
    return i == n;
}

// XDI 映射
namespace XDI {
    static const uint16_t Temperature      = 0x0810;
    static const uint16_t PacketCounter    = 0x1020;
    static const uint16_t SampleTimeCoarse = 0x1070;
    static const uint16_t SampleTimeFine   = 0x1060;
    static const uint16_t Quaternion       = 0x2010;
    static const uint16_t EulerAngles      = 0x2030;
    static const uint16_t Acceleration     = 0x4020;
    static const uint16_t FreeAcceleration = 0x4030;
    static const uint16_t RateOfTurn       = 0x8020;
    static const uint16_t MagneticField    = 0xC020;
    static const uint16_t StatusWord       = 0xE020;
    static const uint16_t DeviceId         = 0xE080;
    static const uint16_t LocationId       = 0xE081;
}

// IMU数据帧结构
struct ImuFrameData {
    bool has_counter = false; 
    uint32_t counter = 0;
    bool has_tfine = false; 
    uint32_t t_fine = 0;
    bool has_tcoarse = false; 
    uint32_t t_coarse = 0;
    bool has_quat = false; 
    float qw = 0, qx = 0, qy = 0, qz = 0;
    bool has_euler = false; 
    float roll = 0, pitch = 0, yaw = 0;
    bool has_acc = false; 
    float ax = 0, ay = 0, az = 0;
    bool has_facc = false; 
    float fax = 0, fay = 0, faz = 0;
    bool has_gyr = false; 
    float gx = 0, gy = 0, gz = 0;
    bool has_mag = false;
    float mx = 0, my = 0, mz = 0;
    bool has_status = false; 
    uint32_t status = 0;
    bool has_temperature = false;
    float temperature = 0.0f;

    void reset() { 
        *this = ImuFrameData(); 
    }
};

class ImuPublisherNode : public rclcpp::Node {
public:
    ImuPublisherNode() : Node("imu_publisher") {
        // 声明参数
        this->declare_parameter("device", "/dev/ttyS0");
        this->declare_parameter("baudrate", 115200);
        this->declare_parameter("frame_id", "imu_link");
        this->declare_parameter("publish_rate", 100.0);
        this->declare_parameter("verbose", false);

        // 获取参数
        device_ = this->get_parameter("device").as_string();
        baudrate_ = this->get_parameter("baudrate").as_int();
        frame_id_ = this->get_parameter("frame_id").as_string();
        publish_rate_ = this->get_parameter("publish_rate").as_double();
        verbose_ = this->get_parameter("verbose").as_bool();

        // 创建发布者
        imu_pub_ = this->create_publisher<sensor_msgs::msg::Imu>("/imu/data", 10);
        mag_pub_ = this->create_publisher<sensor_msgs::msg::MagneticField>("/imu/mag", 10);

        // 初始化IMU传感器
        sensor_ = std::make_unique<ImuSensor>(device_, baudrate_);
        if (!sensor_->init()) {
            RCLCPP_ERROR(this->get_logger(), "串口初始化失败: %s", device_.c_str());
            return;
        }

        RCLCPP_INFO(this->get_logger(), "读取设备: %s, 波特率: %d", device_.c_str(), baudrate_);
        RCLCPP_INFO(this->get_logger(), "发布频率: %.1f Hz", publish_rate_);

        // 创建定时器进行非阻塞读取
        auto timer_period = std::chrono::milliseconds(static_cast<int>(1000.0 / publish_rate_));
        timer_ = this->create_wall_timer(timer_period, std::bind(&ImuPublisherNode::timer_callback, this));

        // 预留缓冲区
        stream_.reserve(8192);
        payload_.reserve(1024);
    }

private:
    void timer_callback() {
        read_and_process_data();
    }

    void read_and_process_data() {
        uint8_t tmp[512];
        ssize_t n = sensor_->readSome(tmp, sizeof(tmp));
        
        if (n < 0) {
            RCLCPP_ERROR(this->get_logger(), "读取数据失败");
            return;
        } else if (n == 0) {
            // 非阻塞模式下没有数据可读，正常情况
            return;
        }

        if (verbose_) {
            RCLCPP_INFO(this->get_logger(), "读取 %zd 字节数据", n);
        }

        stream_.insert(stream_.end(), tmp, tmp + n);

        // 解析 XBUS 帧
        size_t i = 0;
        while (i < stream_.size()) {
            // 寻找 PRE=0xFA
            while (i < stream_.size() && stream_[i] != 0xFA) ++i;
            if (i + 3 >= stream_.size()) break; // 需要 BusID, MID, LEN
            
            size_t start = i;
            uint8_t pre = stream_[i++];
            (void)pre;
            uint8_t busId = stream_[i++];
            uint8_t mid = stream_[i++];
            uint8_t len = stream_[i++];

            size_t need = 0;
            int extBytes = 0;
            if (len == 0xFF) {
                if (i + 2 > stream_.size()) { i = start + 1; break; }
                extBytes = 2;
                need = (size_t(stream_[i]) << 8) | stream_[i+1];
                i += 2;
            } else {
                need = len;
            }

            if (need > (1<<16)) { // 长度过大，丢弃该起点
                i = start + 1;
                continue;
            }

            if (i + need + 1 > stream_.size()) { // payload + CS 未齐
                i = start; 
                break;
            }

            const uint8_t* pl = &stream_[i];
            uint8_t cs = stream_[i + need];
            uint8_t calc = checksum_basic(busId, mid, pl, need, extBytes);

            if (cs != calc) {
                if (verbose_) {
                    RCLCPP_WARN(this->get_logger(), "校验失败 bus=0x%02X mid=0x%02X need=%zu got_cs=0x%02X calc=0x%02X", 
                               busId, mid, need, cs, calc);
                }
                i = start + 1;
                continue;
            }

            payload_.assign(pl, pl + need);
            i += need + 1;

            // 仅处理 MTData2
            if (mid == 0x36) {
                ImuFrameData frame_data;
                bool ok = foreachPacket(payload_, [&](uint16_t did, const uint8_t* p, uint8_t sz){
                    if (did == XDI::PacketCounter && (sz==2 || sz==4)) { 
                        frame_data.has_counter = true; 
                        frame_data.counter = (sz==2)? be16(p) : be32(p); 
                    }
                    else if (did == XDI::SampleTimeFine && sz==4) { 
                        frame_data.has_tfine = true; 
                        frame_data.t_fine = be32(p); 
                    }
                    else if (did == XDI::SampleTimeCoarse && sz==4) { 
                        frame_data.has_tcoarse = true; 
                        frame_data.t_coarse = be32(p); 
                    }
                    else if (did == XDI::Quaternion && sz==16) { 
                        frame_data.has_quat = true; 
                        frame_data.qw = befloat32(p); 
                        frame_data.qx = befloat32(p+4); 
                        frame_data.qy = befloat32(p+8); 
                        frame_data.qz = befloat32(p+12); 
                    }
                    else if (did == XDI::EulerAngles && sz==12) { 
                        frame_data.has_euler = true; 
                        frame_data.roll = befloat32(p); 
                        frame_data.pitch = befloat32(p+4); 
                        frame_data.yaw = befloat32(p+8); 
                    }
                    else if (did == XDI::Acceleration && sz==12) {
                        frame_data.has_acc = true; 
                        frame_data.ax = befloat32(p); 
                        frame_data.ay = befloat32(p+4); 
                        frame_data.az = befloat32(p+8);
                    }
                    else if (did == XDI::FreeAcceleration && sz==12) {
                        frame_data.has_facc = true; 
                        frame_data.fax = befloat32(p); 
                        frame_data.fay = befloat32(p+4); 
                        frame_data.faz = befloat32(p+8);
                    }
                    else if (did == XDI::RateOfTurn && sz==12) { 
                        frame_data.has_gyr = true; 
                        frame_data.gx = befloat32(p); 
                        frame_data.gy = befloat32(p+4); 
                        frame_data.gz = befloat32(p+8); 
                    }
                    else if (did == XDI::MagneticField && sz==12) {
                        frame_data.has_mag = true;
                        frame_data.mx = befloat32(p);
                        frame_data.my = befloat32(p+4);
                        frame_data.mz = befloat32(p+8);
                    }
                    else if (did == XDI::StatusWord && sz==4) { 
                        frame_data.has_status = true; 
                        frame_data.status = be32(p); 
                    }
                    else if (did == XDI::Temperature && sz==4) {
                        frame_data.has_temperature = true;
                        frame_data.temperature = befloat32(p);
                    }
                });

                if (ok) {
                    publish_imu_data(frame_data);
                } else if (verbose_) {
                    RCLCPP_WARN(this->get_logger(), "MTData2 包截断，丢弃本帧");
                }
            } else if (verbose_) {
                RCLCPP_INFO(this->get_logger(), "忽略 MID=0x%02X, payload=%zu 字节", mid, need);
            }
        }

        // 清理已消费数据
        if (i > 0) {
            stream_.erase(stream_.begin(), stream_.begin() + i);
        }
    }

    void publish_imu_data(const ImuFrameData& data) {
        auto now = this->get_clock()->now();

        // 发布IMU数据
        auto imu_msg = sensor_msgs::msg::Imu();
        imu_msg.header.stamp = now;
        imu_msg.header.frame_id = frame_id_;

        // 线加速度 (使用自由加速度，如果没有则使用普通加速度)
        if (data.has_facc) {
            imu_msg.linear_acceleration.x = data.fax;
            imu_msg.linear_acceleration.y = data.fay;
            imu_msg.linear_acceleration.z = data.faz;
        } else if (data.has_acc) {
            imu_msg.linear_acceleration.x = data.ax;
            imu_msg.linear_acceleration.y = data.ay;
            imu_msg.linear_acceleration.z = data.az;
        }

        // 角速度
        if (data.has_gyr) {
            imu_msg.angular_velocity.x = data.gx;
            imu_msg.angular_velocity.y = data.gy;
            imu_msg.angular_velocity.z = data.gz;
        }

        // 四元数
        if (data.has_quat) {
            imu_msg.orientation.w = data.qw;
            imu_msg.orientation.x = data.qx;
            imu_msg.orientation.y = data.qy;
            imu_msg.orientation.z = data.qz;
        }

        imu_pub_->publish(imu_msg);

        // 发布磁场数据
        if (data.has_mag) {
            auto mag_msg = sensor_msgs::msg::MagneticField();
            mag_msg.header.stamp = now;
            mag_msg.header.frame_id = frame_id_;
            mag_msg.magnetic_field.x = data.mx;
            mag_msg.magnetic_field.y = data.my;
            mag_msg.magnetic_field.z = data.mz;
            mag_pub_->publish(mag_msg);
        }

        if (verbose_) {
            RCLCPP_INFO(this->get_logger(), 
                "发布IMU数据 - 加速度: [%.3f, %.3f, %.3f], 角速度: [%.3f, %.3f, %.3f], 四元数: [%.3f, %.3f, %.3f, %.3f]",
                imu_msg.linear_acceleration.x, imu_msg.linear_acceleration.y, imu_msg.linear_acceleration.z,
                imu_msg.angular_velocity.x, imu_msg.angular_velocity.y, imu_msg.angular_velocity.z,
                imu_msg.orientation.w, imu_msg.orientation.x, imu_msg.orientation.y, imu_msg.orientation.z);
        }
    }

    // ROS2相关
    rclcpp::Publisher<sensor_msgs::msg::Imu>::SharedPtr imu_pub_;
    rclcpp::Publisher<sensor_msgs::msg::MagneticField>::SharedPtr mag_pub_;
    rclcpp::TimerBase::SharedPtr timer_;

    // IMU传感器相关
    std::unique_ptr<ImuSensor> sensor_;
    std::string device_;
    int baudrate_;
    std::string frame_id_;
    double publish_rate_;
    bool verbose_;

    // 数据缓冲区
    std::vector<uint8_t> stream_;
    std::vector<uint8_t> payload_;
};

int main(int argc, char** argv) {
    rclcpp::init(argc, argv);
    
    try {
        auto node = std::make_shared<ImuPublisherNode>();
        rclcpp::spin(node);
    } catch (const std::exception& e) {
        std::cerr << "错误: " << e.what() << std::endl;
        return EXIT_FAILURE;
    }
    
    rclcpp::shutdown();
    return EXIT_SUCCESS;
}
